{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11403553",
   "metadata": {},
   "outputs": [],
   "source": [
    "# V3 Imports\n",
    "from sagemaker.train import ModelTrainer\n",
    "from sagemaker.train.configs import Compute, SourceCode, InputData, StoppingCondition\n",
    "from sagemaker.train.tuner import HyperparameterTuner\n",
    "from sagemaker.core.parameter import ContinuousParameter, CategoricalParameter\n",
    "from sagemaker.core.helper.session_helper import get_execution_role\n",
    "from sagemaker.mlops.workflow.steps import TuningStep\n",
    "from sagemaker.mlops.workflow.model_step import ModelStep\n",
    "from sagemaker.serve.model_builder import ModelBuilder\n",
    "from sagemaker.core.workflow.pipeline_context import PipelineSession"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2af3e16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize SageMaker session\n",
    "pipeline_session = PipelineSession()\n",
    "region = pipeline_session.boto_region_name\n",
    "default_bucket = pipeline_session.default_bucket()\n",
    "\n",
    "# Role Configuration\n",
    "# Option 1: Auto-detect (works in SageMaker Studio/Notebook instances)\n",
    "# Option 2: Manually specify your SageMaker execution role ARN\n",
    "try:\n",
    "    role = get_execution_role()\n",
    "    print(f\"✓ Auto-detected role: {role}\")\n",
    "except Exception as e:\n",
    "    print(f\"⚠️  Could not auto-detect role: {e}\")\n",
    "    # Manually specify your SageMaker execution role ARN here:\n",
    "    role = \"<IAM Role ARN>\"\n",
    "    print(f\"✓ Using manually specified role: {role}\")\n",
    "\n",
    "# Define prefixes for organization\n",
    "prefix = \"v3-tuning\"\n",
    "base_job_prefix = \"pytorch-mnist-hpo\"\n",
    "\n",
    "# Configuration\n",
    "training_instance_type = \"ml.m5.xlarge\"\n",
    "account_id = pipeline_session.account_id()\n",
    "local_dir = \"data\"\n",
    "\n",
    "print(f\"\\nRegion: {region}\")\n",
    "print(f\"Role: {role}\")\n",
    "print(f\"Bucket: {default_bucket}\")\n",
    "print(f\"Prefix: {prefix}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65e4f1ce",
   "metadata": {},
   "source": [
    "### Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54552596",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download MNIST dataset\n",
    "from torchvision.datasets import MNIST\n",
    "from torchvision import transforms\n",
    "\n",
    "MNIST.mirrors = [\n",
    "    f\"https://sagemaker-example-files-prod-{region}.s3.amazonaws.com/datasets/image/MNIST/\"\n",
    "]\n",
    "\n",
    "print(\"Downloading MNIST dataset...\")\n",
    "MNIST(\n",
    "    local_dir,\n",
    "    download=True,\n",
    "    transform=transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,))\n",
    "    ]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99e03d7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload to S3\n",
    "s3_data_uri = pipeline_session.upload_data(\n",
    "    path=local_dir,\n",
    "    bucket=default_bucket,\n",
    "    key_prefix=f\"{prefix}/data\"\n",
    ")\n",
    "\n",
    "print(f\"Training data uploaded to: {s3_data_uri}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0e8e423",
   "metadata": {},
   "source": [
    "### Tune Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "620c2c66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure source code\n",
    "source_code = SourceCode(\n",
    "    source_dir=\".\",  # Current directory containing mnist.py\n",
    "    entry_script=\"mnist.py\"\n",
    ")\n",
    "\n",
    "# Configure compute resources\n",
    "compute = Compute(\n",
    "    instance_type=training_instance_type,\n",
    "    instance_count=1,\n",
    "    volume_size_in_gb=30\n",
    ")\n",
    "\n",
    "# Configure stopping condition\n",
    "stopping_condition = StoppingCondition(\n",
    "    max_runtime_in_seconds=3600  # 1 hour\n",
    ")\n",
    "\n",
    "# Get PyTorch training image\n",
    "training_image = f\"763104351884.dkr.ecr.{region}.amazonaws.com/pytorch-training:1.10.0-gpu-py38\"\n",
    "\n",
    "# Create ModelTrainer\n",
    "model_trainer = ModelTrainer(\n",
    "    training_image=training_image,\n",
    "    source_code=source_code,\n",
    "    compute=compute,\n",
    "    stopping_condition=stopping_condition,\n",
    "    hyperparameters={\n",
    "        \"epochs\": 1,  # Use 1 epoch for faster tuning\n",
    "        \"backend\": \"gloo\"\n",
    "    },\n",
    "    sagemaker_session=pipeline_session,\n",
    "    role=role,\n",
    "    base_job_name=base_job_prefix\n",
    ")\n",
    "\n",
    "print(\"ModelTrainer configured successfully\")\n",
    "print(f\"Training Image: {training_image}\")\n",
    "print(f\"Instance Type: {training_instance_type}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0247410f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define hyperparameter ranges to tune\n",
    "hyperparameter_ranges = {\n",
    "    \"lr\": ContinuousParameter(0.001, 0.1),\n",
    "    \"batch-size\": CategoricalParameter([32, 64, 128, 256, 512]),\n",
    "}\n",
    "\n",
    "# Define objective metric\n",
    "objective_metric_name = \"average test loss\"\n",
    "objective_type = \"Minimize\"\n",
    "\n",
    "# Define metric definitions\n",
    "metric_definitions = [\n",
    "    {\n",
    "        \"Name\": \"average test loss\",\n",
    "        \"Regex\": \"Test set: Average loss: ([0-9\\\\.]+)\"\n",
    "    }\n",
    "]\n",
    "\n",
    "# Create HyperparameterTuner\n",
    "tuner = HyperparameterTuner(\n",
    "    model_trainer=model_trainer,\n",
    "    objective_metric_name=objective_metric_name,\n",
    "    hyperparameter_ranges=hyperparameter_ranges,\n",
    "    metric_definitions=metric_definitions,\n",
    "    max_jobs=3,\n",
    "    max_parallel_jobs=2,\n",
    "    strategy=\"Random\",\n",
    "    objective_type=objective_type,\n",
    "    early_stopping_type=\"Auto\"\n",
    ")\n",
    "\n",
    "print(\"HyperparameterTuner configured successfully\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa51514e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare input data\n",
    "training_data = InputData(\n",
    "    channel_name=\"training\",\n",
    "    data_source=s3_data_uri\n",
    ")\n",
    "\n",
    "# Start tuning job\n",
    "print(\"Starting hyperparameter tuning job...\")\n",
    "tuner_run_args = tuner.tune(\n",
    "    inputs=[training_data],\n",
    "    wait=False\n",
    ")\n",
    "\n",
    "step_tuning = TuningStep(\n",
    "    name=\"HPTuning\",\n",
    "    step_args=tuner_run_args,\n",
    ")\n",
    "\n",
    "\n",
    "tuning_job_name = tuner._current_job_name\n",
    "print(f\"\\nTuning job started: {tuning_job_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43163a10",
   "metadata": {},
   "source": [
    "### Deploy best tuned model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92f74453",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "model_builder = ModelBuilder(\n",
    "    image_uri=training_image,\n",
    "    s3_model_data_url=step_tuning.get_top_model_s3_uri(\n",
    "        top_k=0, s3_bucket=default_bucket, prefix=base_job_prefix\n",
    "    ),\n",
    "    sagemaker_session=pipeline_session,\n",
    "    role_arn=role,\n",
    ")\n",
    "\n",
    "step_create_best = ModelStep(\n",
    "    name=\"CreateBestModel\",\n",
    "    step_args=model_builder.build(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "164b2881",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.mlops.workflow.pipeline import Pipeline\n",
    "\n",
    "pipeline = Pipeline(\n",
    "    name=\"pipeline-v3\",\n",
    "    steps=[step_tuning, step_create_best],\n",
    "    sagemaker_session=pipeline_session,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4db9ba7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This step is slow because source directory will be uploaded it to S3.\n",
    "pipeline.definition()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0731be8",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline.upsert(role_arn=role)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8619aaa6",
   "metadata": {},
   "outputs": [],
   "source": [
    "execution = pipeline.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88a6364b",
   "metadata": {},
   "outputs": [],
   "source": [
    "execution.describe()['PipelineExecutionStatus']"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
